package com.bjy.qa.agent.mock.http;

import org.apache.commons.io.output.ByteArrayOutputStream;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;

/**
 * 缓存读取 HttpServletRequest 的 body 的包装器（解决多次读取 body 报错问题）
 */
public class BufferedHttpServletRequestWrapper extends HttpServletRequestWrapper {
    private byte[] buffer; // 缓存 body

    /**
     * 构造方法，读取 request body 并缓存
     * @param request
     */
    public BufferedHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
        try {
            InputStream is = request.getInputStream();
            ByteArrayOutputStream os = new ByteArrayOutputStream();
            byte[] tmp = new byte[1024];
            int read;
            while ((read = is.read(tmp)) > 0) {
                os.write(tmp, 0, read);
            }
            this.buffer = os.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 返回缓存的 body
     * @return
     */
    @Override
    public ServletInputStream getInputStream() {
        return new BufferedServletInputStream(this.buffer);
    }

    /**
     * 缓存的 ServletInputStream
     */
    class BufferedServletInputStream extends ServletInputStream {
        private ByteArrayInputStream inputStream;

        public BufferedServletInputStream(byte[] buffer) {
            this.inputStream = new ByteArrayInputStream(buffer);
        }

        @Override
        public int available() {
            return inputStream.available();
        }

        @Override
        public int read() throws IOException {
            return inputStream.read();
        }

        @Override
        public boolean isFinished() {
            return inputStream.available() == 0;
        }

        @Override
        public boolean isReady() {
            return true;
        }

        @Override
        public void setReadListener(ReadListener listener) {
            throw new RuntimeException("Not implemented");
        }
    }
}
